{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "ce118a4b-3c3f-42f6-ae1d-8b825fbccb93",
            "metadata": {},
            "source": [
                "# Newton Raphson Method In Two Dimensions"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5e12d3fa-68b7-4c43-81c5-515b07a0b33d",
            "metadata": {},
            "source": [
                "## Learning Outcomes\n",
                "This example teaches how to compute the solution for systems of equations in two variables using NumPy. There are two equations, $f_{1}(x,y)$ and $f_{2}(x, y)$, with two variables each, $x$ and $y$. We seek to find a solution that satisfies these two equations using Newton's method. To understand Newton's method in multiple dimensions, please see [this](https://wiki.math.ntnu.no/_media/tma4125/2017v/newton.pdf) note by Markus Grasmair.\n",
                "\n",
                "The example also teaches how to interpret a warning from cuPyNumeric when the import statement is changed from importing numpy to importing cuPyNumeric.\n",
                "\n",
                "---"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a814eecb-682b-4573-a2c8-572e5f0638f7",
            "metadata": {},
            "source": [
                "## Background\n",
                "We consider the following functions,\n",
                "\n",
                "$$\n",
                "f_{1}(x,y) = x^{2} + y^{2} - 13 = 0\n",
                "$$\n",
                "\n",
                "$$\n",
                "f_{2}(x,y) = x^{2} - 2y^{2} + 14 = 0\n",
                "$$\n",
                "\n",
                "and their Jacobian, $J$, \n",
                "\n",
                "$$\n",
                "J = \\begin{bmatrix}\n",
                " \\frac{\\partial f_{1}}{\\partial x} & \\frac{\\partial f_{1}}{\\partial y} \\\\\n",
                " \\frac{\\partial f_{2}}{\\partial x} & \\frac{\\partial f_{2}}{\\partial y}\n",
                "\\end{bmatrix}\n",
                "$$\n",
                "\n",
                "\n",
                "Substituting the functions, $f_{1}(x, y)$ and $f_{2}(x, y)$, we get,\n",
                "\n",
                "$$\n",
                "J = \\begin{matrix}\n",
                " 2x & 2y \\\\\n",
                " 2x & -4y\n",
                "\\end{matrix}\n",
                "$$"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "795b4478-cd32-438d-b806-a1215f3a07bb",
            "metadata": {},
            "source": [
                "## Implementation"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "id": "7a0284e7-fe55-4137-95a8-0fff06d535bf",
            "metadata": {},
            "outputs": [],
            "source": [
                "import numpy as np"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "id": "4252e070-a4d4-4b07-87ad-c03e789f062a",
            "metadata": {},
            "outputs": [],
            "source": [
                "def function(x: np.ndarray) -> np.ndarray:\n",
                "    \"Return a numpy array that has the computed values of $f_{1}(x, y)$ and $f_{2}(x, y)$\"\n",
                "    return np.array([np.sum(x**2) - 13.0, x[0] ** 2 - 2.0 * x[1] ** 2 + 14.0])\n",
                "\n",
                "\n",
                "def jacobian(x: np.ndarray) -> np.ndarray:\n",
                "    \"Return a 2x2 numpy array that has the computed values of the Jacobian, J\"\n",
                "    return np.array([[2 * x[0], 2 * x[1]], [2.0 * x[0], -4.0 * x[1]]])"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "136e1298-9af2-4fc3-9b1a-a56afbd54d7a",
            "metadata": {},
            "source": [
                "Setup an iterative loop that updates an initial guess $x_{k} =  x_{k-1} - {[\\mathbf{J}(x_{k})]}^{-1} \\cdot \\mathbf{f}(x_{k})$\\\n",
                "To compute the inverse of the matrix, $\\mathbf{J}$, we use the `inv` API from NumPy's `linalg` package, and to determine when to terminate the loop, \\\n",
                "we compute the L2 norm of the difference in solution between two iterations and check if it is less than a specified tolerance."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a91752f1-5ca8-44dd-9a26-525cdf87ab51",
            "metadata": {},
            "source": [
                "When you switch the import statement from importing to importing cupynumeric, you might see a warning like this:\n",
                "\n",
                "---\n",
                "\n",
                "*RuntimeWarning: cuPyNumeric has not implemented inv and is falling back to canonical NumPy. You may notice significantly decreased performance for this function call.*\n",
                "\n",
                "---\n",
                "\n",
                "This means that cuPyNumeric has not implemented the `linalg.inv` API and is falling back to NumPy's implementation. This means that the API would be *eagerly* executed using NumPy's single-threaded implementation. If the API was intended to be invoked from a GPU, the data will get transferred from the GPU to the CPU before the API is executed. This can have performance implications, as indicated by the warning."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "id": "c243f28e-ad5e-4c64-8340-96922785c253",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "\n",
                        "Newton's method converged in 7 iterations to xk: [-2.  3.]\n"
                    ]
                }
            ],
            "source": [
                "# number of iterations to try\n",
                "niters = 20\n",
                "\n",
                "# tolerance that sets the accuracy of solution\n",
                "tol = 1e-6\n",
                "\n",
                "# print additional information\n",
                "verbose = False\n",
                "\n",
                "# initial guess\n",
                "xk = np.array([-20.0, 20.0])\n",
                "\n",
                "# Newton's method\n",
                "for iter in range(niters):\n",
                "    xk_old = xk\n",
                "\n",
                "    if verbose:\n",
                "        print(f\"iter: {iter}, xk: {xk}\")\n",
                "    xk = xk - np.linalg.inv(jacobian(xk)).dot(function(xk))\n",
                "\n",
                "    l2_norm = np.linalg.norm((xk - xk_old))\n",
                "    if l2_norm < tol:\n",
                "        break\n",
                "\n",
                "# let the user know if the solution converged or not\n",
                "if iter == niters - 1:\n",
                "    print(\n",
                "        f\"\\nNewton's method did not converge for this function, tolerance ({tol}) and number of iterations ({niters})\"\n",
                "    )\n",
                "else:\n",
                "    print(f\"\\nNewton's method converged in {iter} iterations to xk: {xk}\")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "e5a2e401-e058-4bcc-ac0c-4caa80102079",
            "metadata": {},
            "source": [
                "---\n",
                "\n",
                "We see that the solution has converged to $(x, y) = (-2, 3)$ which satisfies both the equation in 7 iterations\n",
                "\n",
                "The problem can be cast such that the computation of inverse is substituted by a linear solve, as shown below:\\\n",
                "$x_{k} =  x_{k-1} - x_{k}^{*}$\\\n",
                "$x_{k}^{*} = {[\\mathbf{J}(x_{k})]}^{-1} \\cdot \\mathbf{f}(x_{k})$\n",
                "\n",
                "And $x_{k}^{*} $ is solution to the system of equation defined as ${\\mathbf{J}(x_{k})}~ x_{k}^{*} = \\mathbf{f}(x_{k})$\n",
                "\n",
                "---\n",
                "\n",
                "We can then use NumPy's `linalg.solve` API to perform the linear solve as shown below. And we can see that the algorithm converges to the same solution in exactly the same number of iteration"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "id": "11527885-0be6-4ebf-80fa-9dec85bb0c3c",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "\n",
                        "Newton's method converged in 7 iterations to xk: [-2.  3.]\n"
                    ]
                }
            ],
            "source": [
                "# number of iterations to try\n",
                "niters = 20\n",
                "\n",
                "# tolerance that sets the accuracy of solution\n",
                "tol = 1e-6\n",
                "\n",
                "# print additional information\n",
                "verbose = False\n",
                "\n",
                "# initial guess\n",
                "xk = np.array([-20.0, 20.0])\n",
                "\n",
                "# Newton's method\n",
                "for iter in range(niters):\n",
                "    xk_old = xk\n",
                "\n",
                "    if verbose:\n",
                "        print(f\"iter: {iter}, xk: {xk}\")\n",
                "    xk = xk - np.linalg.solve(\n",
                "        jacobian(xk), function(xk)\n",
                "    )  ## This uses linalg.solve\n",
                "\n",
                "    l2_norm = np.linalg.norm((xk - xk_old))\n",
                "    if l2_norm < tol:\n",
                "        break\n",
                "\n",
                "# let the user know if the solution converged or not\n",
                "if iter == niters - 1:\n",
                "    print(\n",
                "        f\"\\nNewton's method did not converge for this function, tolerance ({tol}) and number of iterations ({niters})\"\n",
                "    )\n",
                "else:\n",
                "    print(f\"\\nNewton's method converged in {iter} iterations to xk: {xk}\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "0c9f494c-518a-4f78-9e88-1aeb2221fa1b",
            "metadata": {},
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.10.14"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}
